import numpy as np
import pandas as pd
from script_convert_rating import filter_unused_data


def convert_ratings(input_fname, output_fname):
    rating = pd.read_csv(input_fname, names=['userID', 'itemID', 'rating', 'timestamp', 'price'], sep='\t', skiprows=1)
    # rating.rename(columns={'movieId': 'itemID'}, inplace=True)
    # rating.rename(columns={'userId': 'userID'}, inplace=True)
    rating.drop(columns=['timestamp', 'price'], inplace=True)

    rating = filter_unused_data.convert_ratings(rating)

    # order_l = ['userID', 'itemID', 'rating']
    # rating = rating[order_l]
    rating.to_csv(output_fname, index=False)
    print("complete {}".format(output_fname))


def convert_ratings_sample(input_fname, output_fname, n_sample_user, n_sample_item):
    rating = pd.read_csv(input_fname, names=['userID', 'itemID', 'rating', 'timestamp', 'price'], sep='\t', skiprows=1)
    # rating.rename(columns={'movieId': 'itemID'}, inplace=True)
    # rating.rename(columns={'userId': 'userID'}, inplace=True)
    rating.drop(columns=['timestamp', 'price'], inplace=True)

    rating = filter_unused_data.convert_ratings(rating)

    n_user = len(np.unique(rating['userID']))
    n_item = len(np.unique(rating['itemID']))

    print(
        f"n_total_user {n_user}, n_total_item {n_item}, n_sample_user {n_sample_user}, n_sample_item {n_sample_item}")

    # sample_userID_l = np.random.choice(np.arange(n_user), n_sample_user) + 1
    # sample_itemID_l = np.random.choice(np.arange(n_item), n_sample_item) + 1
    # rating = rating[
    #     (rating['userID'].isin(sample_userID_l)) & (rating['itemID'].isin(sample_itemID_l))]
    # rating = filter_unused_data.convert_ratings(rating)

    # sample_itemID_l = np.random.choice(np.arange(n_item), n_sample_item) + 1
    sample_itemID_l = np.argsort(np.bincount(rating['itemID']))[:n_sample_item]
    rating = rating[(rating['itemID'].isin(sample_itemID_l+1))]
    rating = filter_unused_data.convert_ratings(rating)

    # rating = rating.iloc[:n_sample_user * n_sample_item // 5000]
    # rating = filter_unused_data.convert_ratings(rating)

    print(np.max(np.bincount(rating['userID'])), np.max(np.bincount(rating['itemID'])))
    n_actual_user = len(np.unique(rating['userID']))
    n_actual_item = len(np.unique(rating['itemID']))
    print(
        f"n_actual_user {n_actual_user}, n_actual_item {n_actual_item}")

    # order_l = ['userID', 'itemID', 'rating']
    # rating = rating[order_l]
    rating.to_csv(output_fname, index=False)
    print("complete {}".format(output_fname))
